// ImGui Win32 + DirectX9 binding
// In this binding, ImTextureID is used to store a 'LPDIRECT3DTEXTURE9' texture identifier. Read the FAQ about ImTextureID in imgui.cpp.

// You can copy and use unmodified imgui_impl_* files in your project. See main.cpp for an example of using this.
// If you use this binding you'll need to call 4 functions: ImGui_ImplXXXX_Init(), ImGui_ImplXXXX_NewFrame(), ImGui::Render() and ImGui_ImplXXXX_Shutdown().
// If you are new to ImGui, see examples/README.txt and documentation at the top of imgui.cpp.
// https://github.com/ocornut/imgui

#include "..\imgui.h"
#include "imgui_dx9.h"

// DirectX
#include <d3dx9.h>
#define DIRECTINPUT_VERSION 0x0800
#include <dinput.h>

// Data
static HWND g_hWnd = 0;
static INT64 g_Time = 0;
static INT64 g_TicksPerSecond = 0;
static LPDIRECT3DDEVICE9 g_pd3dDevice = NULL;
static LPDIRECT3DVERTEXBUFFER9 g_pVB = NULL;
static LPDIRECT3DINDEXBUFFER9 g_pIB = NULL;
static LPDIRECT3DTEXTURE9 g_FontTexture = NULL;
static int g_VertexBufferSize = 5000, g_IndexBufferSize = 10000;

struct CUSTOMVERTEX
{
	D3DXVECTOR3 pos;
	D3DCOLOR col;
	D3DXVECTOR2 uv;
};

#define D3DFVF_CUSTOMVERTEX (D3DFVF_XYZ|D3DFVF_DIFFUSE|D3DFVF_TEX1)

// This is the main rendering function that you have to implement and provide to ImGui (via setting up 'RenderDrawListsFn' in the ImGuiIO structure)
// If text or lines are blurry when integrating ImGui in your engine:
// - in your Render function, try translating your projection matrix by (0.5f,0.5f) or (0.375f,0.375f)
void ImGui_ImplDX9_RenderDrawLists( ImDrawData* draw_data )
{
	// Create and grow buffers if needed
	if( !g_pVB || g_VertexBufferSize < draw_data->TotalVtxCount )
	{
		if( g_pVB )
		{
			g_pVB->Release();
			g_pVB = NULL;
		}
		g_VertexBufferSize = draw_data->TotalVtxCount + 5000;
		if( g_pd3dDevice->CreateVertexBuffer( g_VertexBufferSize * sizeof(CUSTOMVERTEX), D3DUSAGE_DYNAMIC | D3DUSAGE_WRITEONLY, D3DFVF_CUSTOMVERTEX, D3DPOOL_DEFAULT, &g_pVB, NULL ) < 0 )
			return;
	}
	if( !g_pIB || g_IndexBufferSize < draw_data->TotalIdxCount )
	{
		if( g_pIB )
		{
			g_pIB->Release();
			g_pIB = NULL;
		}
		g_IndexBufferSize = draw_data->TotalIdxCount + 10000;
		if( g_pd3dDevice->CreateIndexBuffer( g_IndexBufferSize * sizeof(ImDrawIdx), D3DUSAGE_DYNAMIC | D3DUSAGE_WRITEONLY, sizeof(ImDrawIdx) == 2 ? D3DFMT_INDEX16 : D3DFMT_INDEX32, D3DPOOL_DEFAULT, &g_pIB, NULL ) < 0 )
			return;
	}

	// Backup some DX9 state (not all!)
	// FIXME: Backup/restore everything else
	D3DRENDERSTATETYPE last_render_state_to_backup[] = { D3DRS_CULLMODE, D3DRS_LIGHTING, D3DRS_ZENABLE, D3DRS_ALPHABLENDENABLE, D3DRS_ALPHATESTENABLE, D3DRS_BLENDOP, D3DRS_SRCBLEND, D3DRS_DESTBLEND, D3DRS_SCISSORTESTENABLE };
	DWORD last_render_state_values[ARRAYSIZE(last_render_state_to_backup)] = { 0 };
	IDirect3DPixelShader9* last_ps;
	g_pd3dDevice->GetPixelShader( &last_ps );
	IDirect3DVertexShader9* last_vs;
	g_pd3dDevice->GetVertexShader( &last_vs );
	for( int n = 0; n < ARRAYSIZE(last_render_state_to_backup); n++ )
		g_pd3dDevice->GetRenderState( last_render_state_to_backup[ n ], &last_render_state_values[ n ] );

	// Copy and convert all vertices into a single contiguous buffer
	CUSTOMVERTEX* vtx_dst;
	ImDrawIdx* idx_dst;
	if( g_pVB->Lock( 0, ( UINT )( draw_data->TotalVtxCount * sizeof(CUSTOMVERTEX) ), ( void** )&vtx_dst, D3DLOCK_DISCARD ) < 0 )
		return;
	if( g_pIB->Lock( 0, ( UINT )( draw_data->TotalIdxCount * sizeof(ImDrawIdx) ), ( void** )&idx_dst, D3DLOCK_DISCARD ) < 0 )
		return;
	for( int n = 0; n < draw_data->CmdListsCount; n++ )
	{
		const ImDrawList* cmd_list = draw_data->CmdLists[ n ];
		const ImDrawVert* vtx_src = &cmd_list->VtxBuffer[ 0 ];
		for( int i = 0; i < cmd_list->VtxBuffer.size(); i++ )
		{
			vtx_dst->pos.x = vtx_src->pos.x;
			vtx_dst->pos.y = vtx_src->pos.y;
			vtx_dst->pos.z = 0.0f;
			vtx_dst->col = ( vtx_src->col & 0xFF00FF00 ) | ( ( vtx_src->col & 0xFF0000 ) >> 16 ) | ( ( vtx_src->col & 0xFF ) << 16 ); // RGBA --> ARGB for DirectX9
			vtx_dst->uv.x = vtx_src->uv.x;
			vtx_dst->uv.y = vtx_src->uv.y;
			vtx_dst++;
			vtx_src++;
		}
		memcpy( idx_dst, &cmd_list->IdxBuffer[ 0 ], cmd_list->IdxBuffer.size() * sizeof(ImDrawIdx) );
		idx_dst += cmd_list->IdxBuffer.size();
	}
	g_pVB->Unlock();
	g_pIB->Unlock();
	g_pd3dDevice->SetStreamSource( 0, g_pVB, 0, sizeof( CUSTOMVERTEX) );
	g_pd3dDevice->SetIndices( g_pIB );
	g_pd3dDevice->SetFVF( D3DFVF_CUSTOMVERTEX );

	// Setup render state: fixed-pipeline, alpha-blending, no face culling, no depth testing
	g_pd3dDevice->SetPixelShader( NULL );
	g_pd3dDevice->SetVertexShader( NULL );
	g_pd3dDevice->SetRenderState( D3DRS_CULLMODE, D3DCULL_NONE );
	g_pd3dDevice->SetRenderState( D3DRS_LIGHTING, false );
	g_pd3dDevice->SetRenderState( D3DRS_ZENABLE, false );
	g_pd3dDevice->SetRenderState( D3DRS_ALPHABLENDENABLE, true );
	g_pd3dDevice->SetRenderState( D3DRS_ALPHATESTENABLE, false );
	g_pd3dDevice->SetRenderState( D3DRS_BLENDOP, D3DBLENDOP_ADD );
	g_pd3dDevice->SetRenderState( D3DRS_SRCBLEND, D3DBLEND_SRCALPHA );
	g_pd3dDevice->SetRenderState( D3DRS_DESTBLEND, D3DBLEND_INVSRCALPHA );
	g_pd3dDevice->SetRenderState( D3DRS_SCISSORTESTENABLE, true );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_COLOROP, D3DTOP_MODULATE );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_COLORARG1, D3DTA_TEXTURE );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_COLORARG2, D3DTA_DIFFUSE );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_ALPHAOP, D3DTOP_MODULATE );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_ALPHAARG1, D3DTA_TEXTURE );
	g_pd3dDevice->SetTextureStageState( 0, D3DTSS_ALPHAARG2, D3DTA_DIFFUSE );
	g_pd3dDevice->SetSamplerState( 0, D3DSAMP_MINFILTER, D3DTEXF_LINEAR );
	g_pd3dDevice->SetSamplerState( 0, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR );

	// Setup orthographic projection matrix
	D3DXMATRIXA16 mat;
	D3DXMatrixIdentity( &mat );
	g_pd3dDevice->SetTransform( D3DTS_WORLD, &mat );
	g_pd3dDevice->SetTransform( D3DTS_VIEW, &mat );
	D3DXMatrixOrthoOffCenterLH( &mat, 0.5f, ImGui::GetIO().DisplaySize.x + 0.5f, ImGui::GetIO().DisplaySize.y + 0.5f, 0.5f, -1.0f, +1.0f );
	g_pd3dDevice->SetTransform( D3DTS_PROJECTION, &mat );

	// Render command lists
	int vtx_offset = 0;
	int idx_offset = 0;
	for( int n = 0; n < draw_data->CmdListsCount; n++ )
	{
		const ImDrawList* cmd_list = draw_data->CmdLists[ n ];
		for( int cmd_i = 0; cmd_i < cmd_list->CmdBuffer.size(); cmd_i++ )
		{
			const ImDrawCmd* pcmd = &cmd_list->CmdBuffer[ cmd_i ];
			if( pcmd->UserCallback )
			{
				pcmd->UserCallback( cmd_list, pcmd );
			}
			else
			{
				const RECT r = { ( LONG )pcmd->ClipRect.x, ( LONG )pcmd->ClipRect.y, ( LONG )pcmd->ClipRect.z, ( LONG )pcmd->ClipRect.w };
				g_pd3dDevice->SetTexture( 0, ( LPDIRECT3DTEXTURE9 )pcmd->TextureId );
				g_pd3dDevice->SetScissorRect( &r );
				g_pd3dDevice->DrawIndexedPrimitive( D3DPT_TRIANGLELIST, vtx_offset, 0, ( UINT )cmd_list->VtxBuffer.size(), idx_offset, pcmd->ElemCount / 3 );
			}
			idx_offset += pcmd->ElemCount;
		}
		vtx_offset += cmd_list->VtxBuffer.size();
	}

	// Restore some modified DX9 state (not all!)
	g_pd3dDevice->SetPixelShader( last_ps );
	g_pd3dDevice->SetVertexShader( last_vs );
	for( int n = 0; n < ARRAYSIZE(last_render_state_to_backup); n++ )
		g_pd3dDevice->SetRenderState( last_render_state_to_backup[ n ], last_render_state_values[ n ] );
}

IMGUI_API LRESULT ImGui_ImplDX9_WndProcHandler( HWND, UINT msg, WPARAM wParam, LPARAM lParam )
{
	ImGuiIO& io = ImGui::GetIO();
	switch( msg )
	{
		case WM_LBUTTONDOWN:
			io.MouseDown[ 0 ] = true;
			return true;
		case WM_LBUTTONUP:
			io.MouseDown[ 0 ] = false;
			return true;
		case WM_RBUTTONDOWN:
			io.MouseDown[ 1 ] = true;
			return true;
		case WM_RBUTTONUP:
			io.MouseDown[ 1 ] = false;
			return true;
		case WM_MBUTTONDOWN:
			io.MouseDown[ 2 ] = true;
			return true;
		case WM_MBUTTONUP:
			io.MouseDown[ 2 ] = false;
			return true;
		case WM_MOUSEWHEEL:
			io.MouseWheel += GET_WHEEL_DELTA_WPARAM(wParam) > 0 ? +1.0f : -1.0f;
			return true;
		case WM_MOUSEMOVE:
			io.MousePos.x = ( signed short )( lParam );
			io.MousePos.y = ( signed short )( lParam >> 16 );
			return true;
		case WM_KEYDOWN:
			if( wParam < 256 )
				io.KeysDown[ wParam ] = true;
			return true;
		case WM_KEYUP:
			if( wParam < 256 )
				io.KeysDown[ wParam ] = false;
			return true;
		case WM_CHAR:
			// You can also use ToAscii()+GetKeyboardState() to retrieve characters.
			if( wParam > 0 && wParam < 0x10000 )
				io.AddInputCharacter( ( unsigned short )wParam );
			return true;
	}
	return 0;
}

bool ImGui_ImplDX9_Init( void* hwnd, IDirect3DDevice9* device )
{
	g_hWnd = ( HWND )hwnd;
	g_pd3dDevice = device;

	if( !QueryPerformanceFrequency( ( LARGE_INTEGER * )&g_TicksPerSecond ) )
		return false;
	if( !QueryPerformanceCounter( ( LARGE_INTEGER * )&g_Time ) )
		return false;

	ImGuiIO& io = ImGui::GetIO();
	io.KeyMap[ ImGuiKey_Tab ] = VK_TAB; // Keyboard mapping. ImGui will use those indices to peek into the io.KeyDown[] array that we will update during the application lifetime.
	io.KeyMap[ ImGuiKey_LeftArrow ] = VK_LEFT;
	io.KeyMap[ ImGuiKey_RightArrow ] = VK_RIGHT;
	io.KeyMap[ ImGuiKey_UpArrow ] = VK_UP;
	io.KeyMap[ ImGuiKey_DownArrow ] = VK_DOWN;
	io.KeyMap[ ImGuiKey_PageUp ] = VK_PRIOR;
	io.KeyMap[ ImGuiKey_PageDown ] = VK_NEXT;
	io.KeyMap[ ImGuiKey_Home ] = VK_HOME;
	io.KeyMap[ ImGuiKey_End ] = VK_END;
	io.KeyMap[ ImGuiKey_Delete ] = VK_DELETE;
	io.KeyMap[ ImGuiKey_Backspace ] = VK_BACK;
	io.KeyMap[ ImGuiKey_Enter ] = VK_RETURN;
	io.KeyMap[ ImGuiKey_Escape ] = VK_ESCAPE;
	io.KeyMap[ ImGuiKey_A ] = 'A';
	io.KeyMap[ ImGuiKey_C ] = 'C';
	io.KeyMap[ ImGuiKey_V ] = 'V';
	io.KeyMap[ ImGuiKey_X ] = 'X';
	io.KeyMap[ ImGuiKey_Y ] = 'Y';
	io.KeyMap[ ImGuiKey_Z ] = 'Z';

	io.RenderDrawListsFn = ImGui_ImplDX9_RenderDrawLists; // Alternatively you can set this to NULL and call ImGui::GetDrawData() after ImGui::Render() to get the same ImDrawData pointer.
	io.ImeWindowHandle = g_hWnd;

	return true;
}

void ImGui_ImplDX9_Shutdown()
{
	ImGui_ImplDX9_InvalidateDeviceObjects();
	ImGui::Shutdown();
	g_pd3dDevice = NULL;
	g_hWnd = 0;
}

static bool ImGui_ImplDX9_CreateFontsTexture()
{
	// Build texture atlas
	ImGuiIO& io = ImGui::GetIO();
	unsigned char* pixels;
	int width, height, bytes_per_pixel;
	io.Fonts->GetTexDataAsRGBA32( &pixels, &width, &height, &bytes_per_pixel );

	// Upload texture to graphics system
	g_FontTexture = NULL;
	if( D3DXCreateTexture( g_pd3dDevice, width, height, 1, D3DUSAGE_DYNAMIC, D3DFMT_A8B8G8R8, D3DPOOL_DEFAULT, &g_FontTexture ) < 0 )
		return false;
	D3DLOCKED_RECT tex_locked_rect;
	if( g_FontTexture->LockRect( 0, &tex_locked_rect, NULL, 0 ) != D3D_OK )
		return false;
	for( int y = 0; y < height; y++ )
		memcpy( ( unsigned char * )tex_locked_rect.pBits + tex_locked_rect.Pitch * y, pixels + ( width * bytes_per_pixel ) * y, ( width * bytes_per_pixel ) );
	g_FontTexture->UnlockRect( 0 );

	// Store our identifier
	io.Fonts->TexID = ( void * )g_FontTexture;

	return true;
}

bool ImGui_ImplDX9_CreateDeviceObjects()
{
	if( !g_pd3dDevice )
		return false;
	if( !ImGui_ImplDX9_CreateFontsTexture() )
		return false;
	return true;
}

void ImGui_ImplDX9_InvalidateDeviceObjects()
{
	if( !g_pd3dDevice )
		return;
	if( g_pVB )
	{
		g_pVB->Release();
		g_pVB = NULL;
	}
	if( g_pIB )
	{
		g_pIB->Release();
		g_pIB = NULL;
	}
	if( LPDIRECT3DTEXTURE9 tex = ( LPDIRECT3DTEXTURE9 )ImGui::GetIO().Fonts->TexID )
	{
		tex->Release();
		ImGui::GetIO().Fonts->TexID = 0;
	}
	g_FontTexture = NULL;
}

void ImGui_ImplDX9_NewFrame()
{
	if( !g_FontTexture )
		ImGui_ImplDX9_CreateDeviceObjects();

	ImGuiIO& io = ImGui::GetIO();

	// Setup display size (every frame to accommodate for window resizing)
	RECT rect;
	GetClientRect( g_hWnd, &rect );
	io.DisplaySize = ImVec2( ( float )( rect.right - rect.left ), ( float )( rect.bottom - rect.top ) );

	// Setup time step
	INT64 current_time;
	QueryPerformanceCounter( ( LARGE_INTEGER * )&current_time );
	io.DeltaTime = ( float )( current_time - g_Time ) / g_TicksPerSecond;
	g_Time = current_time;

	// Read keyboard modifiers inputs
	io.KeyCtrl = ( GetKeyState( VK_CONTROL ) & 0x8000 ) != 0;
	io.KeyShift = ( GetKeyState( VK_SHIFT ) & 0x8000 ) != 0;
	io.KeyAlt = ( GetKeyState( VK_MENU ) & 0x8000 ) != 0;
	io.KeySuper = false;
	// io.KeysDown : filled by WM_KEYDOWN/WM_KEYUP events
	// io.MousePos : filled by WM_MOUSEMOVE events
	// io.MouseDown : filled by WM_*BUTTON* events
	// io.MouseWheel : filled by WM_MOUSEWHEEL events

	// Hide OS mouse cursor if ImGui is drawing it
	SetCursor( io.MouseDrawCursor ? NULL : LoadCursor( NULL, IDC_ARROW ) );

	// Start the frame
	ImGui::NewFrame();
}
